1324F - Maximum White Subtree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
# define ll long long int
//ll m=998244353;
//__________________fast function_____________
void fast()
{
//___don't use it in interactive problems___
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
}
vector<vector<int>>v;
vector<int>vis,a,ans,ans2;
void fun(int n)
{
    vis[n]=1;
    int val=0;
    if(a[n]==1)
    {
        val++;
    }
    else
    {
        val--;
    }
    for(auto i: v[n])
    {
        if(vis[i]==-1)
        {
            fun(i);
        val+=max(ans[i],0);
        }
    }
    ans[n]=val;
}
void fun2(int par,int n)
{
    vis[n]=1;
    ans2[n]=ans[n];
    if(par>0)
    {
        if(ans[n]>0)
        {
            ans2[n]+=max(ans2[par]-ans[n],0);
        }
        else
        {
            ans2[n]+=max(ans2[par],0);
        }
    }
    for(auto i: v[n])
    {
        if(vis[i]==-1)
        {
            fun2(n,i);
        }
    }
}

int main()
{
    fast();
    int n;
    cin>>n;
    ans.clear();
    ans.resize(n+1);
    a.clear();
    a.resize(n+1);
    ans2.clear();
    ans2.resize(n+1);
    for(int i=1;i<=n;i++)
    {
        cin>>a[i];
    }
    v.clear();
    vis.clear();
    v.resize(n+1);
    vis.resize(n+1,-1);
    for(int i=0;i<n-1;i++)
    {
        int x,y;
        cin>>x>>y;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    fun(1);
    vis.clear();
    vis.resize(n+1,-1);

    fun2(-1,1);
    for(int i=1;i<=n;i++)
    {
        cout<<ans2[i]<<" ";
    }
    cout<<"\n";
    return 0;
}